import torch 
from model import CLSTM, CNN
from torch.utils.data import DataLoader, Dataset
from config import data_dir, batch_size, clip_grad, output_dir, embed_size, feature_size, hidden_size, lr, epoches, \
                src_path, tgt_path, src_path_standard, tgt_path_standard
from scipy import signal
import numpy as np 

def get_right_num(pred, tgt):
    sum_sample = len(pred)
    right = 0
    pred_arg = torch.argmax(pred, dim=1)
    pred_arg = pred_arg.view(-1)
    tgt = tgt.view(-1)
    for index in range(sum_sample):
        if int(pred_arg[index].item()) == int(tgt[index].item()):
            right += 1

    return right 

def get_accuracy(pred, tgt):
    sum_sample = len(pred)
    right = 0
    pred_arg = torch.argmax(pred, dim=1)
    pred_arg = pred_arg.view(-1)
    tgt = tgt.view(-1)
    for index in range(sum_sample):
        if int(pred_arg[index].item()) == int(tgt[index].item()):
            right += 1

    return float(right) / sum_sample 

class EegDataset(Dataset):
    def __init__(self, src_data, tgt_data):
        super(EegDataset).__init__()
        self.src_data = src_data
        self.tgt_data = tgt_data
        wn = 2 * 80 / 1000
        self.b, self.a = signal.butter(6, wn, "lowpass")
        
    def __getitem__(self, i):
        src, tgt = self.src_data[i], self.tgt_data[i]
        src = src - torch.mean(src, dim=0)
        detrend_data = signal.detrend(src, axis=0)
        filted_data = signal.filtfilt(self.b, self.a, detrend_data, axis=0)
        data = filted_data.astype(np.float32).copy()
        data = torch.tensor(data)

        return data, tgt

    def __len__(self):
        return len(self.src_data)

if __name__ == "__main__":
    # model_path = "./model_dict.pkl"
    # state_dict = torch.load(model_path, map_location="cpu")

    # model = SportClassifierModel(embed_size, hidden_size, feature_size, "cpu")
    # model.load_state_dict(state_dict)
    # src = torch.load("src_data_mean_29_30.pkl")
    # tgt = torch.load("tgt_data.pkl")
    # print(src[:5])
    # print(torch.mean(src))
    # print(src[:5] * 10)
    # print(tgt[:5])
    # print(model(src[:5] * 10))

    # model_path = "./model_dict_feature.pkl"
    # state_dict = torch.load(model_path, map_location="cpu")

    # model = SimpleClassifier(embed_size, hidden_size, feature_size, "cpu")
    # model.load_state_dict(state_dict)
    # model.eval()
    # src = torch.load("psd_feature_data.pkl")
    # tgt = torch.load("psd_tgt_data.pkl")

    # print(get_accuracy(model(src[1500:]), tgt[1500:]))

    # print(get_accuracy(model(src[:1500]), tgt[:1500]))


    model_path = "./model_dict_cnn_11.pkl"
    
    state_dict = torch.load(model_path, map_location="cpu")
    model = CNN("cpu")
    model.load_state_dict(state_dict)
    model.eval()
    src = torch.load("./s3_test_src.pkl")
    tgt = torch.load("./s3_test_tgt.pkl")
    dataset = EegDataset(src, tgt)
    dataloader = DataLoader(dataset, batch_size=10, shuffle=True)

    right_num = 0
    for test_src, test_tgt in dataloader:
        right_num += get_right_num(model(test_src), test_tgt)
        print(model(test_src))
        print(test_tgt)
    print(right_num)

    # print(get_accuracy(model(src[:100]), tgt[:100]))
    


    # pred = torch.tensor([[0.5, 0.7], [0.6, 0.8], [0.9, 0.1]])
    # tgt = torch.tensor([1, 1, 1])
    # print(get_accuracy(pred, tgt))


    


    # src_data = torch.load("./src_data_0628.pkl")
    # tgt_data = torch.load("./tgt_data_0628.pkl")
    # print(tgt_data[:5])

    # test_src =  src_data[1000:]
    # test_tgt = tgt_data[1000:]

    # print(test_src.shape)




    




